from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import pickle
import json
import numpy as np
import cv2
DATA_PATH = '../../data/wider/'
DEBUG = False
import os


'''
#Values    Name      Description
----------------------------------------------------------------------------
   1    type         Describes the type of object: 'Face'
   4    bbox         2D face bounding box (x, y, w, h).
   1    quality      Float from 0 (low quality) to 1 (high quality), face quality
   15   landmark     5 Facial landmarks of each face (x1, y1, v1, x2, y2, v2, ..., x5, y5, v5)
   1    score        Only for results: Float, indicating confidence in
                     detection, needed for p/r curves, higher is better.
'''

cats = ['Face']
cat_info = []
for i, cat in enumerate(cats):
  cat_info.append({'name': cat, 'id': i + 1})

for split in ['train', 'val', 'test']:
  ret = {'images': [], 'annotations': [], "categories": cat_info}
  path_to_annos = {}
  path = None 
  for line in open(DATA_PATH + 'labels/{}.txt'.format(split), 'r'):
    line = line.strip()
    if line.startswith('#'):
      path = line[1:].strip()
      path_to_annos[path] = []
    else:
      assert path in path_to_annos, path
      path_to_annos[path].append(line)

  for path in path_to_annos:
    im = cv2.imread(DATA_PATH + 'images/{}/{}'.format(split, path))
    height, width = im.shape[:2]
    image_id = len(ret['images']) + 1
    image_info = {'file_name': path, 'id': image_id, 'height': height, 'width': width}
    ret['images'].append(image_info)
    for line in path_to_annos[path]:
      tmp = list(map(float, line.split()))
      annotation_id = len(ret['annotations']) + 1
      if min(tmp[2:4]) <= 0:
        print(tmp[:4])
        continue
      ann = {'image_id': image_id,
              'id': annotation_id,
              'category_id': 1,
              'bbox': tmp[:4],
              'iscrowd': len(path_to_annos[path]) > 10,
              'area': tmp[2] * tmp[3]}
      if split == 'train':
        ann['quality'] = tmp[19]
        ann['landmark'] = tmp[4:19]
      ret['annotations'].append(ann)
      if DEBUG:
        im = cv2.rectangle(im, (int(ann['bbox'][0]), int(ann['bbox'][1])),
                           (int(ann['bbox'][0]) + int(ann['bbox'][2]), 
                            int(ann['bbox'][1]) + int(ann['bbox'][3])),
                           (255, 0, 0), 2)
        if split == 'train':
          im = cv2.putText(im, str(ann['quality']), (int(ann['bbox'][0]), int(ann['bbox'][1])-10),                 cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255))
          for i in range(5):
            im = cv2.circle(im, (int(ann['landmark'][3*i]), int(ann['landmark'][3*i+1])), 1, (0,255,0), 2)
        print(ann)
    if DEBUG:
      im = cv2.putText(im, str(image_id), (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 2, (0,0,255))
      cv2.imshow('image', im)
      cv2.waitKey()

  print("# images: ", len(ret['images']))
  print("# annotations: ", len(ret['annotations']))
  out_dir = DATA_PATH + 'annotations/'
  if not os.path.exists(out_dir):
    os.makedirs(out_dir)
  json.dump(ret, open(out_dir + 'wider_{}.json'.format(split), 'w'))
